# -*- coding:utf-8 -*-
from cassandra.cluster import Cluster
from cassandra import ConsistencyLevel
from cassandra.query import BatchStatement, BatchType
from datetime import *
import json
import os
from wmcloud_dataclient import Client

DEFAULT_DEDATE = '1970-01-01'
def get_stock_basic_info(raw_msg):
        try:
            stock_basic_info = json.loads(raw_msg)
            if stock_basic_info['retCode'] == 1 and stock_basic_info['retMsg'] == 'Success':
                stock_info_d_list = stock_basic_info['data']
                cluster = Cluster(protocol_version=4)
                session = cluster.connect()
                session.set_keyspace('wmcloud')
                insert_sock = session.prepare('insert into wmcloud.sz_sh_stock_basic_info(ticker ,listdate ,delistdate ,enddate ,totalshares ,equtype ,equtypecd ,exchangecd ,excountrycd ,listsector ,listsectorcd ,liststatuscd ,nonrestfloata ,nonrestfloatshares ,officeaddr ,partyid ,primeoperating ,secfullname ,secid ,secshortname ,transcurrcd ,tshequity) values (? ,? ,? ,? ,? ,? ,? ,? ,? ,? ,? ,? ,? ,? ,? ,? ,? ,? ,? ,? ,? ,?)')
                for data in stock_info_d_list:
                    listdate_timestamp = datetime.strptime(data['listDate'], '%Y-%m-%d') + timedelta(hours=16)
                    delistdate_timestamp = datetime.strptime(data['delistDate'], '%Y-%m-%d') + timedelta(hours=16) if 'delistDate' in data else datetime.strptime(DEFAULT_DEDATE, '%Y-%m-%d')
                    enddate_timestamp = datetime.strptime(data['endDate'], '%Y-%m-%d') + timedelta(hours=16) if 'endDate' in data else datetime.strptime(DEFAULT_DEDATE, '%Y-%m-%d')
                    print(data['ticker'])
                    session.execute(insert_sock, (data['ticker'], listdate_timestamp, delistdate_timestamp, enddate_timestamp,data['totalShares'], data['equType'], data['equTypeCD'], data['exchangeCD'], data['exCountryCD'], data['ListSector'], data['ListSectorCD'], data['listStatusCD'], data['nonrestfloatA'], data['nonrestFloatShares'], data['officeAddr'] if 'officeAddr' in data else 'Unknown', data['partyID'], data['primeOperating'] if 'primeOperating' in data else 'Unknown', data['secFullName'], data['secID'], data['secShortName'], data['transCurrCD'], data['TShEquity'] if 'TShEquity' in data else 0.0))
                session.shutdown()
                cluster.shutdown()
        except Exception as e:
            print(e)

if __name__ == '__main__':
    client = Client('93ef5ebba60c5f1da190801750612423db818daad66155b6b191417b19d331b9')
    url_list = ['/api/equity/getEqu.json?field=&listStatusCD=&secID=&ticker=&equTypeCD=A', '/api/equity/getEqu.json?field=&listStatusCD=&secID=&ticker=&equTypeCD=B']
    for url in url_list:
        try:
            code, result = client.getData(url)
            if code == 200:
                get_stock_basic_info(result)
        except Exception as e:
            print(e)
